import pandas as pd
import numpy as np
from datetime import datetime
import os


# 1. 数据输入与预处理
def load_data(file_path="附件1.xlsx"):
    """
    从Excel文件中读取数据，假设包含门店、商品和顾客需求等列
    返回DataFrame，包含'store' (门店), 'product' (商品), 'demand' (顾客需求)
    """
    try:
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"警告: 文件 '{file_path}' 不存在，将使用示例数据")
            return create_sample_data()

        # 尝试读取Excel文件
        # 假设Excel文件结构：列名分别为 '公司名称' (store), '货号' (product), '顾客需求' (demand)
        df = pd.read_excel(file_path, sheet_name=0, usecols=['公司名称', '货号', '顾客需求'])
        df.columns = ['store', 'product', 'demand']

        # 检查是否有缺失值
        if df['demand'].isnull().any():
            print("警告: 数据中存在缺失值，将用0填充")
            df['demand'] = df['demand'].fillna(0)

        return df
    except Exception as e:
        print(f"读取文件时出错: {e}")
        print("将使用示例数据继续运行")
        return create_sample_data()


def create_sample_data():
    """创建示例数据"""
    sample_data = {
        'store': ['s1', 's1', 's2', 's2', 's3', 's3', 's4', 's4'],
        'product': ['p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2'],
        'demand': [100, 150, 200, 250, 300, 350, 400, 450]
    }
    return pd.DataFrame(sample_data)


# 2. 每日需求预测
def predict_daily_demand(df):
    """
    基于顾客需求数据预测每日需求
    输入：DataFrame 包含门店、商品和需求
    输出：字典 { (store, product): daily_demand }
    """
    daily_demand = {}
    for (store, product), group in df.groupby(['store', 'product']):
        # 取每个门店-商品组合的最大需求作为预测值
        demand_value = group['demand'].max()
        daily_demand[(store, product)] = demand_value
    return daily_demand


# 3. 10分钟粒度需求分配
def allocate_10min_demand(daily_demand):
    """
    将每日需求均匀分配到144个10分钟时段
    输入：每日需求字典
    输出：三维字典 {store, product, t}: demand_per_10min
    """
    demand_10min = {}
    total_intervals = 144  # 一天24小时 × 6个10分钟
    for (store, product), daily_demand_value in daily_demand.items():
        demand_per_10min = daily_demand_value / total_intervals
        for t in range(1, total_intervals + 1):
            demand_10min[(store, product, t)] = round(demand_per_10min, 4)  # 保留4位小数
    return demand_10min


# 4. 主程序
def main():
    # 加载数据
    file_path = "附件1.xlsx"  # 请替换为实际文件路径
    print(f"尝试加载数据文件: {file_path}")
    data = load_data(file_path)

    # 打印输入数据以验证
    print("\n输入数据预览：")
    print(data.head())
    print(f"\n数据总量: {len(data)} 条记录")
    print(f"门店数量: {data['store'].nunique()}")
    print(f"商品种类: {data['product'].nunique()}")

    # 预测每日需求
    daily_demand = predict_daily_demand(data)
    print("\n每日需求预测：")
    for (store, product), demand in list(daily_demand.items())[:5]:  # 只显示前5个
        print(f"门店 {store}, 商品 {product}: {demand} 件")
    print(f"...共 {len(daily_demand)} 个门店-商品组合")

    # 分配10分钟粒度需求
    demand_10min = allocate_10min_demand(daily_demand)
    print("\n10分钟粒度需求分配示例：")
    # 显示不同门店和商品的示例
    samples = [
        (s, p, t) for (s, p, t) in demand_10min.keys()
        if s == 's1' and p == 'p1' and t in [1, 72, 144]  # 第1、72和144时段
    ]
    for key in samples:
        store, product, t = key
        demand = demand_10min[key]
        print(f"门店 {store}, 商品 {product}, 时段 {t}: {demand:.4f} 件")

    # 保存结果到文件
    result_df = pd.DataFrame([
        {'store': store, 'product': product, 'time_slot': t, 'demand': demand}
        for (store, product, t), demand in demand_10min.items()
    ])

    output_file = "demand_10min_output.csv"
    result_df.to_csv(output_file, index=False)
    print(f"\n预测结果已保存到 '{output_file}'")
    print(f"总记录数: {len(result_df)}")

    # 当前时间戳
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"\n程序运行完成，时间：{current_time}")


if __name__ == "__main__":
    print("多门店需求预测系统 v1.0")
    print("=" * 50)
    main()